#!/usr/bin/env python3
"""
Memory Guard Module for Jetson Nano 4GB Memory Protection
内存保护模块，专门针对4GB内存限制和递归代码保护
"""

import psutil
import threading
import time
import gc
import sys
import traceback
from typing import Dict, Any, Optional, Callable
from functools import wraps
import logging


class MemoryGuard:
    """内存保护类，提供内存监控和保护机制"""
    
    def __init__(self, 
                 memory_threshold: float = 95.0,  # 内存使用率阈值(%)
                 swap_threshold: float = 70.0,    # swap使用率阈值(%)
                 max_recursion_depth: int = 50,   # 最大递归深度
                 check_interval: float = 5.0,     # 检查间隔(秒)
                 logger: Optional[logging.Logger] = None):
        """
        初始化内存保护器
        
        Args:
            memory_threshold: 内存使用率警告阈值
            swap_threshold: swap使用率警告阈值
            max_recursion_depth: 最大允许递归深度
            check_interval: 内存检查间隔
            logger: 日志记录器
        """
        self.memory_threshold = memory_threshold
        self.swap_threshold = swap_threshold
        self.max_recursion_depth = max_recursion_depth
        self.check_interval = check_interval
        self.logger = logger or logging.getLogger(__name__)
        
        # 内存监控状态
        self.monitoring = False
        self.monitor_thread = None
        self.last_warning_time = 0
        self.warning_cooldown = 30  # 警告冷却时间(秒)
        
        # 递归深度跟踪
        self.recursion_counters = {}
        self.recursion_lock = threading.Lock()
        
        # 内存统计
        self.memory_stats = {
            'peak_memory_usage': 0.0,
            'peak_swap_usage': 0.0,
            'warning_count': 0,
            'cleanup_count': 0
        }
    
    def get_memory_info(self) -> Dict[str, Any]:
        """获取当前内存信息"""
        try:
            # 获取内存信息
            memory = psutil.virtual_memory()
            swap = psutil.swap_memory()
            
            memory_usage_percent = memory.percent
            swap_usage_percent = swap.percent if swap.total > 0 else 0.0
            
            # 更新峰值统计
            self.memory_stats['peak_memory_usage'] = max(
                self.memory_stats['peak_memory_usage'], 
                memory_usage_percent
            )
            self.memory_stats['peak_swap_usage'] = max(
                self.memory_stats['peak_swap_usage'], 
                swap_usage_percent
            )
            
            return {
                'memory_total_gb': round(memory.total / (1024**3), 2),
                'memory_available_gb': round(memory.available / (1024**3), 2),
                'memory_used_gb': round(memory.used / (1024**3), 2),
                'memory_usage_percent': round(memory_usage_percent, 1),
                'swap_total_gb': round(swap.total / (1024**3), 2),
                'swap_used_gb': round(swap.used / (1024**3), 2),
                'swap_usage_percent': round(swap_usage_percent, 1),
                'memory_critical': memory_usage_percent > self.memory_threshold,
                'swap_critical': swap_usage_percent > self.swap_threshold
            }
        except Exception as e:
            self.logger.error(f"Error getting memory info: {e}")
            return {}
    
    def emergency_cleanup(self) -> bool:
        """紧急内存清理"""
        try:
            self.logger.warning("Performing emergency memory cleanup...")
            
            # 1. 强制垃圾回收
            collected = gc.collect()
            self.logger.info(f"Garbage collection freed {collected} objects")
            
            # 2. 清理Python内部缓存
            sys.intern.clear() if hasattr(sys.intern, 'clear') else None
            
            # 3. 尝试释放未使用的内存
            try:
                import ctypes
                libc = ctypes.CDLL("libc.so.6")
                libc.malloc_trim(0)
            except:
                pass  # 如果失败就忽略
            
            self.memory_stats['cleanup_count'] += 1
            self.logger.info("Emergency cleanup completed")
            return True
            
        except Exception as e:
            self.logger.error(f"Error during emergency cleanup: {e}")
            return False
    
    def check_memory_status(self) -> Dict[str, Any]:
        """检查内存状态并在必要时采取行动"""
        memory_info = self.get_memory_info()
        
        if not memory_info:
            return {'status': 'error', 'message': 'Failed to get memory info'}
        
        current_time = time.time()
        memory_critical = memory_info['memory_critical']
        swap_critical = memory_info['swap_critical']
        
        # 检查是否需要警告
        if (memory_critical or swap_critical) and \
           (current_time - self.last_warning_time) > self.warning_cooldown:
            
            self.memory_stats['warning_count'] += 1
            self.last_warning_time = current_time
            
            warning_msg = f"Memory warning - RAM: {memory_info['memory_usage_percent']}%, " \
                         f"Swap: {memory_info['swap_usage_percent']}%"
            self.logger.warning(warning_msg)
            
            # 如果内存使用率过高，执行紧急清理
            if memory_info['memory_usage_percent'] > 90:
                self.emergency_cleanup()
        
        return {
            'status': 'critical' if (memory_critical and swap_critical) else 
                     'warning' if (memory_critical or swap_critical) else 'normal',
            'memory_info': memory_info,
            'stats': self.memory_stats.copy()
        }
    
    def start_monitoring(self):
        """启动内存监控"""
        if self.monitoring:
            return
        
        self.monitoring = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()
        self.logger.info("Memory monitoring started")
    
    def stop_monitoring(self):
        """停止内存监控"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        self.logger.info("Memory monitoring stopped")
    
    def _monitor_loop(self):
        """内存监控循环"""
        while self.monitoring:
            try:
                self.check_memory_status()
                time.sleep(self.check_interval)
            except Exception as e:
                self.logger.error(f"Error in memory monitor loop: {e}")
                time.sleep(self.check_interval)
    
    def recursion_guard(self, func_name: str = None):
        """递归保护装饰器"""
        def decorator(func: Callable):
            @wraps(func)
            def wrapper(*args, **kwargs):
                # 使用函数名作为递归计数器的键
                counter_key = func_name or f"{func.__module__}.{func.__name__}"
                
                with self.recursion_lock:
                    current_depth = self.recursion_counters.get(counter_key, 0)
                    
                    # 检查递归深度
                    if current_depth >= self.max_recursion_depth:
                        error_msg = f"Maximum recursion depth ({self.max_recursion_depth}) " \
                                   f"exceeded for function {counter_key}"
                        self.logger.error(error_msg)
                        raise RecursionError(error_msg)
                    
                    # 增加递归计数
                    self.recursion_counters[counter_key] = current_depth + 1
                
                try:
                    # 在递归调用前检查内存
                    memory_status = self.check_memory_status()
                    if memory_status['status'] == 'critical':
                        self.logger.warning(f"Critical memory status before {counter_key} execution")
                    
                    # 执行函数
                    result = func(*args, **kwargs)
                    return result
                    
                finally:
                    # 减少递归计数
                    with self.recursion_lock:
                        self.recursion_counters[counter_key] = max(0, 
                            self.recursion_counters.get(counter_key, 1) - 1)
                        
                        # 如果递归深度回到0，清理计数器
                        if self.recursion_counters[counter_key] == 0:
                            del self.recursion_counters[counter_key]
            
            return wrapper
        return decorator
    
    def memory_limit_guard(self, max_memory_percent: float = 90.0):
        """内存限制保护装饰器"""
        def decorator(func: Callable):
            @wraps(func)
            def wrapper(*args, **kwargs):
                # 执行前检查内存
                memory_info = self.get_memory_info()
                if memory_info.get('memory_usage_percent', 0) > max_memory_percent:
                    # 尝试清理内存
                    self.emergency_cleanup()
                    
                    # 再次检查
                    memory_info = self.get_memory_info()
                    if memory_info.get('memory_usage_percent', 0) > max_memory_percent:
                        error_msg = f"Memory usage ({memory_info.get('memory_usage_percent', 0)}%) " \
                                   f"exceeds limit ({max_memory_percent}%) for function {func.__name__}"
                        self.logger.error(error_msg)
                        raise MemoryError(error_msg)
                
                return func(*args, **kwargs)
            return wrapper
        return decorator
    
    def get_stats(self) -> Dict[str, Any]:
        """获取内存保护统计信息"""
        current_memory = self.get_memory_info()
        return {
            'current_memory': current_memory,
            'stats': self.memory_stats.copy(),
            'recursion_counters': self.recursion_counters.copy(),
            'monitoring': self.monitoring,
            'thresholds': {
                'memory_threshold': self.memory_threshold,
                'swap_threshold': self.swap_threshold,
                'max_recursion_depth': self.max_recursion_depth
            }
        }


# 全局内存保护实例
_global_memory_guard = None

def get_memory_guard(logger: Optional[logging.Logger] = None) -> MemoryGuard:
    """获取全局内存保护实例，支持设备化阈值与环境变量配置"""
    global _global_memory_guard
    if _global_memory_guard is None:
        import os
        device_profile = os.environ.get('MUTO_DEVICE_PROFILE', '').lower()
        mem_thresh_env = os.environ.get('MUTO_MEM_THRESHOLD')
        swap_thresh_env = os.environ.get('MUTO_SWAP_THRESHOLD')
        max_rec_env = os.environ.get('MUTO_MAX_RECURSION')
        check_interval_env = os.environ.get('MUTO_CHECK_INTERVAL')

        # 默认阈值
        memory_threshold = 95.0
        swap_threshold = 70.0
        max_recursion_depth = 50
        check_interval = 5.0

        # 基于设备类型的阈值调整
        if device_profile == 'jetson':
            memory_threshold = 90.0
            swap_threshold = 60.0
        elif device_profile == 'raspberry':
            memory_threshold = 92.0
            swap_threshold = 65.0
        elif device_profile == 'pc':
            memory_threshold = 95.0
            swap_threshold = 80.0

        # 环境变量覆盖
        try:
            if mem_thresh_env:
                memory_threshold = float(mem_thresh_env)
            if swap_thresh_env:
                swap_threshold = float(swap_thresh_env)
            if max_rec_env:
                max_recursion_depth = int(max_rec_env)
            if check_interval_env:
                check_interval = float(check_interval_env)
        except Exception:
            pass

        _global_memory_guard = MemoryGuard(
            memory_threshold=memory_threshold,
            swap_threshold=swap_threshold,
            max_recursion_depth=max_recursion_depth,
            check_interval=check_interval,
            logger=logger
        )
        _global_memory_guard.start_monitoring()
    return _global_memory_guard

def memory_protected(func_name: str = None, max_memory_percent: float = 90.0):
    """组合的内存保护装饰器"""
    def decorator(func: Callable):
        guard = get_memory_guard()
        
        # 应用递归保护和内存限制保护
        protected_func = guard.recursion_guard(func_name)(func)
        protected_func = guard.memory_limit_guard(max_memory_percent)(protected_func)
        
        return protected_func
    return decorator


if __name__ == "__main__":
    # 测试代码
    logging.basicConfig(level=logging.INFO)
    
    guard = MemoryGuard()
    guard.start_monitoring()
    
    print("Memory Guard Test")
    print("Current memory status:")
    status = guard.check_memory_status()
    print(f"Status: {status}")
    
    time.sleep(2)
    guard.stop_monitoring()